import logging
import json
import os
from typing import Any, Dict, List, TypedDict, Union, Optional
import queue
import uuid
from ..base.component_base import BaseComponent
from flask import Blueprint, Flask, request
from flask_sock import Sock
import psutil
from .models import WSMessage, WSMessageTypes, WSToServerMessage, PTYMetaDict
import threading
import platform
from .pty import ElementSCIPTY, ElementSCIPTYError
from ...utils import is_windows
from .models import Response

logger = logging.getLogger(__name__)

if platform.system().lower().find("windows") != -1:
    from winpty import PtyProcess, PTY
else:
    import pty

pty_mgr = Blueprint("pty", __name__, url_prefix="/api/pty")


class WSManager:
    def __init__(self) -> None:
        self.websockets = {}

    def add(self, ws):
        self.websockets[id(ws)] = ws

    def remove(self, ws):
        self.websockets.pop(id(ws))


# Responsibility chain pattern.


class WSHandlerCell:
    # 责任链模式，如果无法处理，则抛出KeyError或者NotImplementedError
    def __init__(self) -> None:
        self.type: str = None
        self._next_cell: WSHandlerCell = None

    def handle(self, msg: WSToServerMessage):
        raise NotImplementedError

    def set_next_cell(self, next_cell: "WSHandlerCell"):
        self._next_cell = next_cell

    def handle_in_chain(self, msg: Dict[str, Any]):
        try:
            self.handle(msg)
        except (KeyError, NotImplementedError) as e:
            logger.warning(f"error occurred when handling msg: {e}")
            if self._next_cell is not None:
                self._next_cell.handle(msg)
            else:
                raise NotImplementedError(
                    f"Nextcell was none for handling message {msg}, class {self}!")


class PTYHandleCell(WSHandlerCell):
    def __init__(self) -> None:
        super().__init__()
        self.type = "pty-input"
        self.shell = "cmd" if is_windows() else os.environ.get("SHELL", "sh")
        self.ptys: Dict[str, ElementSCIPTY] = {}

    def get_pty_ids(self) -> List[str]:

        # Get all active ptys
        return [pty_id for pty_id in self.ptys]

    def all_active_ptys(self) -> List[PTYMetaDict]:
        return [self.ptys[pty_id].to_dict() for pty_id in self.ptys]

    def on_pty_close(self, pty_id: str):
        """
        Handle PTY closed event.

        """
        pty = self.ptys.pop(pty_id)
        send_pty_status(pty_id, "closed")

    def close_term(self, termID: str, soft=True):
        """
        Close Terminal

        """
        pty = self.ptys[termID]
        if soft:
            pty.soft_close()
        else:
            raise NotImplementedError

    def handle(self, msg: Dict[str, Any]):
        cmd = msg["cmd"]
        termID = msg["termID"]
        msg_type = msg["msgType"]
        print("handling message", msg)
        assert msg_type in {"new-pty", "close-pty", "cmd-pty"}

        if not termID in self.ptys:
            raise KeyError("Terminal %s not in ptys")
        else:
            self.ptys[termID].write(cmd)

    def create_pty(self, termID: str, cmd: str, name: str) -> Optional[ElementSCIPTY]:
        if termID not in self.ptys:
            self.ptys[termID] = ElementSCIPTY(
                termID, cmd, send_pty_output, self.on_pty_close, name
            )
            return self.ptys[termID]
        else:
            return None

    def resize(self, termID: str, rows: int, cols: int) -> bool:
        if termID not in self.ptys:
            return False
        else:
            self.ptys[termID].resize(rows, cols)
            return True


def send_loop():
    while 1:
        # 如果有若干个相邻的包，就将包全部发送到前端。
        # 这是因为发送包本身有开销，这样可以减少发送包本身的开销。
        values = [send_queue.get()]
        assert isinstance(values[0], WSMessage)

        while not send_queue.empty():
            value = send_queue.get()
            values.append(value)
            assert isinstance(value, WSMessage), value

        ws_keys = list(ws_mgr.websockets.keys())
        for ws_key in ws_keys:
            ws = ws_mgr.websockets[ws_key]
            if ws.connected:
                ws.send(json.dumps([ws_msg.to_dict() for ws_msg in values]))
                # print(f"send at {time.time()}", values[0].payload)
            else:
                ws_mgr.remove(ws)


pty_handle_cell = PTYHandleCell()
ws_handlers = pty_handle_cell


def send_pty_output(term_id: str, content: str):
    if content == "":
        return
    send_queue.put(WSMessage(WSMessageTypes.PTY_OUTPUT, {
                   "output": content, "termID": term_id}))


def send_pty_status(term_id: str, status: str):
    """
    Send PTY status to frontend.

    """
    assert status in {"closed", "opened"}, status
    send_queue.put(
        WSMessage(WSMessageTypes.PTY_STATUS_CHANGE, {
                  "termID": term_id, "status": status})
    )


def send_subprocess_output(type: str, content: str):
    assert type in {"stderr", "stdout"}, "Invalid type" + type
    send_queue.put(WSMessage(WSMessageTypes.SUBPROCESS_OUTPUT,
                             {"type": type, "content": content}))


class TerminalTabComponent(BaseComponent):
    def __init__(self, component_name: str, service: Union[Flask, Blueprint], url_prefix: str = "") -> None:
        super().__init__(component_name, service, url_prefix)

        self.add_post_method("/create", self.new_terminal)
        self.add_get_method("/all", self.get_active_ptys)
        self.add_post_method("/resize", self.resize_terminal)
        self.add_post_method("/close", self.close_terminal)

        self.register_websocket_handlers(service)

        self.send_th = threading.Thread(target=send_loop)
        self.send_th.setDaemon(True)
        self.send_th.start()

    def register_websocket_handlers(self, app: Flask):

        global sock  # 设置成global才可。
        app.config["SOCK_SERVER_OPTIONS"] = {"ping_interval": 25}
        sock = Sock(app)

        @sock.route(f"/api/{self.component_name}/websocket")
        def echo(ws):
            """
            # The ws object has the following methods:
            # - ws.send(data)
            # - ws.receive(timeout=None)
            # - ws.close(reason=None, message=None)
            """
            ws_mgr.add(ws)
            while True:
                data = ws.receive()
                print("received data", data)
                ws_handlers.handle_in_chain(json.loads(data))

    def new_terminal(self, data: dict):
        print(data)
        cmd = data["cmd"]
        name = data["name"]
        new_terminal_id = str(uuid.uuid1())
        try:
            new_term = pty_handle_cell.create_pty(
                new_terminal_id, cmd, name)
            if new_term is None:
                return Response.error("No pty created")
            print(psutil.Process().children(recursive=True))
            return Response.ok(new_term.to_dict())
        except ElementSCIPTYError as e:
            import traceback

            traceback.print_exc()
            return Response.error(str(e))

    def get_active_ptys(self):
        return Response.ok(pty_handle_cell.all_active_ptys())

    def resize_terminal(self, data):
        termID = data["termID"]
        rows = data["rows"]
        cols = data["cols"]
        pty_handle_cell.resize(termID, rows, cols)
        return Response.success_msg("Resize succeeded!")

    def close_terminal(self, data):
        termID = data["termID"]
        pty_handle_cell.close_term(termID)
        return Response.success_msg("Close Succeeded!")


sock = None

send_queue: "queue.Queue[WSMessage]" = queue.Queue(1024)
ws_mgr = WSManager()
